package trace

import (
	"context"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

type (
	Key      = attribute.Key
	KeyValue = attribute.KeyValue
)

type Span struct {
	name        string
	kind        trace.SpanKind
	span        trace.Span
	traceHeader string
	ctx         context.Context
	traceId     string
	spanId      string
}

func StartSpanServer(ctx context.Context, spanName string, traceparent string) *Span {
	tracer := otel.Tracer(traceLibname)

	sp := &Span{
		name: spanName,
		kind: trace.SpanKindServer,
	}

	sp.start(tracer, ctx, traceparent)

	return sp
}

func StartSpanClient(ctx context.Context, spanName string, traceparent string) *Span {
	var tracer trace.Tracer
	if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
		tracer = span.TracerProvider().Tracer(traceLibname)
	} else {
		tracer = otel.Tracer(traceLibname)
	}

	sp := &Span{
		name: spanName,
		kind: trace.SpanKindClient,
	}

	sp.start(tracer, ctx, traceparent)

	return sp
}

func (sp *Span) start(tracer trace.Tracer, ctx context.Context, traceparent string) {
	header := map[string]string{}
	if traceparent != "" {
		header[traceparentHeader] = traceparent
	}

	prop := otel.GetTextMapPropagator()
	ctx = prop.Extract(ctx, TraceContext(header))

	spanctx, span := tracer.Start(ctx, sp.name, trace.WithSpanKind(sp.kind))

	prop.Inject(spanctx, TraceContext(header))

	sp.span = span
	sp.traceHeader = header[traceparentHeader]
	sp.ctx = spanctx

	sctx := trace.SpanContextFromContext(spanctx)

	if sctx.HasTraceID() {
		sp.traceId = sctx.TraceID().String()
	}

	if sctx.HasSpanID() {
		sp.spanId = sctx.SpanID().String()
	}
}

func (sp *Span) End() {
	sp.span.End()
}

func (sp *Span) TraceHeader() string {
	return sp.traceHeader
}

func (sp *Span) Context() context.Context {
	return sp.ctx
}

func (sp *Span) TraceId() string {
	return sp.traceId
}

func (sp *Span) SpanId() string {
	return sp.spanId
}

func (sp *Span) SetAttributes(attrs ...KeyValue) {
	sp.span.SetAttributes(attrs...)
}

func (sp *Span) AddEvent(name string, attrs ...KeyValue) {
	sp.span.AddEvent(name, trace.WithAttributes(attrs...))
}
